# -*- coding:utf-8 -*-
import sys
import struct
import time
import socket
import os
import re
import subprocess

import db_server

DEFAULT_MULTICAST_GROUP = '224.1.1.1'
DEFAULT_MULTICAST_PORT = 1234
DEFAULT_RECEIVER_ADDR = '127.0.0.1'
DEFAULT_RECEIVER_DEV = 'eth0'

class UdpMulticastRecevier():
    def __init__(self):
        self.multicast_group_addr = DEFAULT_MULTICAST_GROUP
        self.receiver_addr = DEFAULT_RECEIVER_ADDR
        self.receiver_port = DEFAULT_MULTICAST_PORT
        self.receiver_dev = DEFAULT_RECEIVER_DEV

class ExecErr(Exception):
    def __init__(self, returncode, cmd, stdout, stderr):
        self.returncode = returncode
        self.cmd = cmd
        self.stdout = stdout
        self.stderr = stderr

def whiptail(args):
    cmd = ['whiptail', '--output-fd', '2']
    cmd.extend(args)
    p = subprocess.Popen(cmd,
                         shell=False,
                         stdin=sys.stdin,
                         stderr=subprocess.PIPE)
    (stdout, stderr) = p.communicate()
    if p.returncode != 0:
        sys.exit()
    return stderr


# 组播地址范围224.0.0.0 到 239.255.255.255 
def get_multicast_group_addr():
    args = ['--title', '请输入组播组地址(默认地址:'+DEFAULT_MULTICAST_GROUP+')', '--inputbox','', '10', '60']
    multicast_group_addr = DEFAULT_MULTICAST_GROUP
    whiptail_args = whiptail(args)
    if whiptail_args != b'':
        multicast_group_addr = str(whiptail_args, encoding='utf-8')
    return multicast_group_addr
# A、B、C类地址
def get_receiver_addr():
    args = ['--title', '请输入接收地址(默认地址:'+DEFAULT_RECEIVER_ADDR+')', '--inputbox','', '10', '60']
    receiver_addr = DEFAULT_RECEIVER_ADDR
    whiptail_args = whiptail(args)
    if whiptail_args != b'':
        receiver_addr = str(whiptail_args, encoding='utf-8')
    return receiver_addr

# 端口范围 1 - 65535
def get_receiver_port():
    args = ['--title', '请输入组播端口(默认端口:'+str(DEFAULT_MULTICAST_PORT)+')', '--inputbox','', '10', '60']
    receiver_port = DEFAULT_MULTICAST_PORT
    whiptail_args = whiptail(args)
    if whiptail_args != b'':
        port = int(whiptail_args)
        if port < 1 or port >65535:
            raise ValueError("The receiver port must be in the range of 1~65535")
        receiver_port = port
    return receiver_port

# 网卡任意填
def get_receiver_dev():
    args = ['--title', '请输入接收网卡(默认网卡:'+DEFAULT_RECEIVER_DEV+')', '--inputbox','', '10', '60']
    receiver_dev = DEFAULT_RECEIVER_DEV
    whiptail_args = whiptail(args)
    if whiptail_args != b'':
        receiver_dev = str(whiptail_args, encoding='utf-8')
    return receiver_dev

def get_receiver():
    receiver = UdpMulticastRecevier()
    receiver.multicast_group_addr = get_multicast_group_addr()
    receiver.receiver_port = get_receiver_port()
    receiver.receiver_addr = get_receiver_addr()
    receiver.receiver_dev = get_receiver_dev()
    return receiver

def receiver():
    receiver = get_receiver()
    print_info(receiver)
    os.system('ip r a '+ receiver.multicast_group_addr +'/32 dev '+ receiver.receiver_dev)
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM, socket.IPPROTO_UDP)
    sock.setsockopt(socket.SOL_SOCKET, 25, bytes(receiver.receiver_dev, encoding='utf-8'))
    sock.bind((receiver.multicast_group_addr, receiver.receiver_port))
    mreq = struct.pack("=4sl", socket.inet_aton(receiver.multicast_group_addr), socket.INADDR_ANY)
    sock.setsockopt(socket.IPPROTO_IP,socket.IP_ADD_MEMBERSHIP,mreq)
    num = 0
    os.system('iptables -w -D INPUT -p udp -d '+ receiver.multicast_group_addr +' --dport '+str(receiver.receiver_port)+' -j ACCEPT')
    os.system('iptables -w -I INPUT 1 -p udp -d '+ receiver.multicast_group_addr +' --dport '+str(receiver.receiver_port)+' -j ACCEPT')
    start = int(time.time())
    traffic = 0
    oldtraffic = 0
    udp_multicast_receiver_id = 0
    while True:
        try:
            message, addr = sock.recvfrom(65535)
            num = num + 1
            if num == 1 :
                if db_server.select_max_receiver_dev() == "unknown":
                    #The sender has started
                    db_server.update_udp_multicast_receiver(receiver.multicast_group_addr, receiver.receiver_addr, receiver.receiver_port, receiver.receiver_dev)
                else:
                    db_server.insert_udp_multicast_receiver(receiver.multicast_group_addr, receiver.receiver_addr, receiver.receiver_port, receiver.receiver_dev)
                udp_multicast_receiver_id = db_server.select_udp_multicast_receiver_id()
            msglen = len(message)
            now = int(time.time())
            if now - start >=1:
                start = now
                pktsbytes = "".join(os.popen('iptables -nvL INPUT -x | grep "'+receiver.multicast_group_addr+'"| grep "udp dpt:'+str(receiver.receiver_port)+'"').readlines())
                pktsbytes = re.sub(' +', ' ', pktsbytes).lstrip()
                pktslist = pktsbytes.split(' ')
                print_info(receiver)
                pktnum = int(pktslist[0])
                traffic = int(pktslist[1])
                traffic_speed = traffic - oldtraffic
                print_real_info(addr, len(message), pktnum, traffic, traffic_speed)
                oldtraffic = traffic
                db_server.insert_udp_multicast_traffic(2, traffic, pktnum, int(time.time()), udp_multicast_receiver_id)
                db_server.insert_udp_multicast_traffic_speed(2, traffic_speed, int(time.time()), udp_multicast_receiver_id)
            if msglen == 5 and message == b'close':
                sock.close()
                pktsbytes = "".join(os.popen('iptables -nvL INPUT -x | grep "'+ receiver.multicast_group_addr +'"| grep "udp dpt:'+str(receiver.receiver_port)+'"').readlines())
                pktsbytes = re.sub(' +', ' ', pktsbytes).lstrip()
                pktslist = pktsbytes.split(' ')
                print_info(receiver)
                pktnum = int(pktslist[0])
                traffic = int(pktslist[1])
                traffic_speed = 0
                print_real_info(addr, len(message), pktnum, traffic, traffic_speed)
                os.system('iptables -w -D INPUT -p udp -d '+ receiver.multicast_group_addr +' --dport '+str(receiver.receiver_port)+' -j ACCEPT')
                os.system('ip r d '+ receiver.multicast_group_addr +'/32 dev '+ receiver.receiver_dev)
                db_server.insert_udp_multicast_traffic(2, traffic, pktnum, int(time.time()), udp_multicast_receiver_id)
                db_server.insert_udp_multicast_traffic_speed(2, traffic_speed, int(time.time()), udp_multicast_receiver_id)
                break
        except :
            if udp_multicast_receiver_id != 0:
                db_server.insert_udp_multicast_traffic(2, traffic, 0, int(time.time()), udp_multicast_receiver_id)
                db_server.insert_udp_multicast_traffic_speed(2, 0, int(time.time()), udp_multicast_receiver_id)
            os.system('iptables -w -D INPUT -p udp -d ' + receiver.multicast_group_addr + ' --dport ' + str(receiver.receiver_port) + ' -j ACCEPT')
            os.system('ip r d ' + receiver.multicast_group_addr + '/32 dev ' + receiver.receiver_dev)
            print("while receive message error occur")
            break

def print_info(receiver):
    os.system('clear')
    print("============================")
    print("  组播地址", receiver.multicast_group_addr)
    print("  组播端口", receiver.receiver_port)
    print("  接收地址", receiver.receiver_addr)
    print("  接收网卡", receiver.receiver_dev)
    print("============================")

def print_real_info(addr, size, pktnum, traffic, traffic_speed):
    print("  发送地址：",addr)
    print("  数据包大小：",size)
    print("  接收包的数量：",pktnum)
    print("  接收的流量：",traffic)
    print("  速率：",traffic_speed)

if __name__ == "__main__":
    receiver()
